import base64
import io

import numpy as np
import torch
from PIL import Image
from flask import Flask, render_template, jsonify

from main00_gan_model_define import Generator  # 从你的模型文件中导入Generator类

app = Flask(__name__)

# 加载生成器模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
z_dim = 100
generator = Generator(z_dim).to(device)
generator.load_state_dict(torch.load('gan_train_cnn/generator.pt', map_location=device))
generator.eval()


@app.route('/')
def index():
    return render_template('index.html')


@app.route('/generate', methods=['POST'])
def generate():
    images = []
    for _ in range(5):
        noise = torch.randn(1, z_dim, device=device)
        with torch.no_grad():
            generated_img = generator(noise).squeeze(0).cpu().numpy()

        # 将生成的图像转换为PIL图像，并进行编码以便在网页中显示
        generated_img = (generated_img * 0.5 + 0.5) * 255  # 反归一化
        generated_img = generated_img.astype(np.uint8)
        img = Image.fromarray(generated_img[0], mode='L')  # 假设是单通道灰度图像
        buffered = io.BytesIO()
        img.save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
        images.append(img_str)

    return jsonify(images=images)


if __name__ == '__main__':
    app.run(debug=True)
